import numpy as np


## 计算dose statistics
def get_dose_statistics_roi_interp(dose_type, binnedDoseValues_, accumulatedRelativeVolumes, voxelVolume, voxelDose, doseAtVolume=0.98):
    """
    Calculates dose statistics within the ROI mask using interpolation.

    Args:
        dose (np.ndarray): 3D dose array.
        roi_mask (np.ndarray): 3D binary mask of the ROI.
        interp_method (str): Interpolation method to use. Default is 'linear'.

    Returns:
        dict: Dictionary containing dose statistics.
    """
    """
    This function calculates dose statistics within the ROI mask using interpolation.
    It takes in a 3D dose array and a 3D binary mask of the ROI, and returns a dictionary
    containing the dose statistics. The interpolation method used can be specified using
    the `interp_method` argument, with the default being 'linear'.
    """
    if dose_type == "average":
        if np.sum(voxelVolume) == 0:
            _msg = "Prescription ROI volume is zero. Cannot scale to average dose."
            # self.logger.error(_msg)
            raise RuntimeError
        else:
            metric = np.sum(voxelDose * voxelVolume) / np.sum(voxelVolume)
    else:
        if dose_type == "D_98":
            dose_at_volume = 0.98
        elif dose_type == "D_95":
            dose_at_volume = 0.95
        elif dose_type == "D_50":
            dose_at_volume = 0.5
        elif dose_type == "D_2":
            dose_at_volume = 0.02
        elif dose_type == "DoseAtVolume":
            dose_at_volume = doseAtVolume
        else:
            _msg = "Prescription type {} not supported. Only 'average', 'D_98', 'D_95', 'D_50', 'D_2' are supported".format(str(dose_type))
            # self.logger.error(_msg)
            raise RuntimeError(_msg)
        metric = get_metric(dose_at_volume, binnedDoseValues_, accumulatedRelativeVolumes)

    return metric


def get_metric(doseAtVolume, binnedDoseValues_, accumulatedRelativeVolumes):
    ind = np.where(accumulatedRelativeVolumes == doseAtVolume)
    if len(ind[0]) > 0:
        assert len(ind[0]) == 1
        metric = binnedDoseValues_[ind]
    else:
        ind = np.where((accumulatedRelativeVolumes - doseAtVolume) < 0)[0][0]
        ind1 = ind - 1
        ind2 = ind
        if accumulatedRelativeVolumes[ind2] != accumulatedRelativeVolumes[ind1]:
            metric = (binnedDoseValues_[ind2] - binnedDoseValues_[ind1]) / (accumulatedRelativeVolumes[ind2] - accumulatedRelativeVolumes[ind1]) * (doseAtVolume - accumulatedRelativeVolumes[ind1]) + binnedDoseValues_[ind1]
        else:
            metric = (binnedDoseValues_[ind2] + binnedDoseValues_[ind1]) * 0.5
    return metric


def accumulate(voxelVolume, voxelDose, doExactAccumulation):
    """
    This function accumulates the dose and prescription values.

    Args:
        voxelVolume (np.ndarray): The voxel volume array.
        voxelDose (np.ndarray): The voxel dose array.
        doExactAccumulation (bool): Whether to perform exact accumulation.

    Returns:
        binnedDoseValues (np.ndarray): .
        accumulatedRelativeVolumes:.
    """
    voxelVolume = voxelVolume.astype("float64")  # 与后端对齐。
    voxelDose = voxelDose.astype("float64")
    voxelVolume = voxelVolume[voxelVolume > 0]  # 过滤出 大于零的
    voxelDose = voxelDose[voxelVolume > 0]  # boolean 索引是 copy
    num = len(voxelDose)
    num2 = np.array([0.0]).astype("float64")
    # if (doExactAccumulation):
    if len(voxelDose) < 5000:  # DVH: Dose and Volume Histogram; 体积-剂量 直方图
        # self.logger.info("voxelDose length {} < 5000, accurate DVH.".format(str(len(voxelDose))))
        accumulatedRelativeVolumes = np.zeros(num + 2)
        binnedDoseValues_ = np.zeros(num + 2)
        # for i in range(0, num + 2):
        #     accumulatedRelativeVolumes[i] = 0.0
        #     binnedDoseValues_[i] = 0.0

        for i in range(0, num):
            accumulatedRelativeVolumes[i + 1] = voxelVolume[i]
            binnedDoseValues_[i + 1] = voxelDose[i]
            if binnedDoseValues_[i + 1] < 0.0:
                binnedDoseValues_[i + 1] = 0.0
        # _dose_vol_array = np.vstack((binnedDoseValues_, accumulatedRelativeVolumes)).transpose()
        # sort_ind = np.argsort(binnedDoseValues_[1:-1])
        # binnedDoseValues_[1:-1] = binnedDoseValues_[1:-1][sort_ind]
        # accumulatedRelativeVolumes[1:-1] = accumulatedRelativeVolumes[1:-1][sort_ind]

        accumulatedRelativeVolumes_middle = accumulatedRelativeVolumes[1:-1]
        binnedDoseValues__middle = binnedDoseValues_[1:-1]
        ind_sort = np.lexsort((accumulatedRelativeVolumes_middle, binnedDoseValues__middle))
        binnedDoseValues_[1:-1] = binnedDoseValues__middle[ind_sort]
        accumulatedRelativeVolumes[1:-1] = accumulatedRelativeVolumes_middle[ind_sort]

    else:
        # self.logger.info("voxelDose length {} >= 5000, coarse DVH.".format(str(len(voxelDose))))
        minMax = [min(voxelDose), max(voxelDose)]
        maxVoxelDose = minMax[1]
        minVoxelDose = max(0.0, minMax[0])
        # 分辨率 1e-6,加1四舍五入; 分割数量最少为2; 最多不超过 1001;
        # 2是 tick 数; 对应 1 个间隔
        splitNum = int(min(1001.0, max(2.0, (maxVoxelDose - minVoxelDose) / 1e-06 + 1.0)))
        assert splitNum > 1

        # 在剂量体积直方图计算中，两个数组一一对应;
        # binnedDoseValues_ 中的每个剂量值，对应 accumulatedRelativeVolumes 中接受该剂量及以上剂量的相对体积
        # [0, min, min+1*delta, min+2delta,..., min+n*delta==max]
        # [>=0, >=min ..., >= max]
        # 累积相对体积; >=某个值的 所有体素的相对体积之和
        accumulatedRelativeVolumes = np.zeros(splitNum + 2)
        # 分组/分箱 - 存储的是按照一定间隔划分的剂量值；数组中的每个元素对应一个剂量区间的代表值
        binnedDoseValues_ = np.zeros(splitNum + 2)

        # for i in range(0, len(accumulatedRelativeVolumes)):
        #     accumulatedRelativeVolumes[i] = 0.0
        #     binnedDoseValues_[i] = 0.0

        if maxVoxelDose > 1e-06:
            # 等差数列; 等价于 binnedDoseValues_[1:splitNum+1] = np.linspace(minVoxelDose, maxVoxelDose, splitNum)
            intervalVoxelDose = (maxVoxelDose - minVoxelDose) / (splitNum - 1)
            for j in range(0, splitNum):
                binnedDoseValues_[j + 1] = minVoxelDose + j * intervalVoxelDose

            for k in range(0, num):
                if voxelDose[k] > minVoxelDose:
                    # 四舍五入的 index; min -> 0; 获取 dose 对应的 bin, left bin(较小的那个阶梯);
                    intervalIndex = (int)(((voxelDose[k] - minVoxelDose) / intervalVoxelDose) + 0.5)
                else:
                    intervalIndex = 0
                    # idx 0 -> dose 0; 因此 intervalIndex + 1, 从 min 开始
                accumulatedRelativeVolumes[intervalIndex + 1] += voxelVolume[k]  # TODO 非常慢

    for i in range(len(accumulatedRelativeVolumes) - 2, 0, -1):
        num9 = accumulatedRelativeVolumes[i]
        accumulatedRelativeVolumes[i] = num2 + 0.5 * num9
        num2 += num9

    # 边界值; 两边取 水平沿
    accumulatedRelativeVolumes[0] = num2
    accumulatedRelativeVolumes[len(accumulatedRelativeVolumes) - 1] = 0.0
    binnedDoseValues_[0] = binnedDoseValues_[1]
    binnedDoseValues_[len(binnedDoseValues_) - 1] = binnedDoseValues_[len(binnedDoseValues_) - 2]
    if not (num2 <= 0.0):
        for l in range(0, len(accumulatedRelativeVolumes)):
            accumulatedRelativeVolumes[l] /= num2
    return np.array(binnedDoseValues_), np.array(accumulatedRelativeVolumes)
